for lr in [0.01, 0.001, 0.0001, 0.00005, 0.00001]:
    a = plot_action_value_by_episode_and_reset(df[(df['model'] == 1) & (df['learning_rate'] == lr) & (df['replay_ratio'] == 16)], 100)
    display(a)

for lr in [0.01, 0.001, 0.0001, 0.00005, 0.00001]:
    a = plot_reward_by_episode_and_reset(df[(df['model'] == 1) & (df['learning_rate'] == lr) & (df['replay_ratio'] == 16)], 100)
    display(a)
    

for lr in [0.01, 0.001, 0.0001, 0.00005, 0.00001]:
    a = plot_hyperparams(df[(df['model'] == 1) & (df['learning_rate'] == lr)], 100)
    display(a)
    

for lr in [0.01, 0.001, 0.0001, 0.00005, 0.00001]:
    a = plot_model_comparison(df[(df['replay_ratio'] == 16) & (df['learning_rate'] == lr)], 100)
    display(a)

for lr in [0.01, 0.001, 0.0001, 0.00005, 0.00001]:
    a = plot_model_compare_overtime(df[(df['learning_rate'] == lr) & (df['weight_reset'] == 10) & (df['replay_ratio'] == 16)], 100)
    display(a)

# Alternative way to see full grid search at once. couldnt figure out how to overlay text
alt.Chart(df[df['model'] == 1].groupby(['weight_reset', 'learning_rate', 'replay_ratio']).mean().reset_index(), width=200, height=150).mark_bar().encode(
    alt.Y('reward:Q', title='Mean Reward', scale=alt.Scale(domain=[0, 0.7])),
    alt.X("weight_reset:N"),
    alt.Color("weight_reset:N").title("Reset Frequency").legend(orient="top", titleOrient="left"),
    alt.Row("learning_rate:Q").title("Learning Rate").header(labelAngle=0),
    alt.Column("replay_ratio:N").title("Replay Ratio"),
)